#include "mex.h"
#include "matrix.h"
#include <cstdio>
#include <cstdlib>
#include <vector> 
#include <time.h>
#include <errno.h>
#include "./source/pnmfile.h"
#include "./source/filter.h"
#include "./source/image.h"
#include "./source/misc.h"
#include "./source/matrix.h"

    //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

using namespace vlib;
using namespace std;

double large = 1E10F;


    //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
/*
 * Simple data model for foreground/background segmentation
 * We assume the foreground is bright and background is dark
 */
void loadData(image<uchar> *I, image<double> * &fg, image<double> * &bg) {
  int width = I->width();
  int height = I->height();
  
  fg = new image<double>(width, height);
  bg = new image<double>(width, height);
  double robustThreshold = 300;//30; dolev
  for (int y = 0; y < height; y++) {
    for (int x = 0; x < width; x++) {
      imRef(fg, x, y) = abs(150-imRef(I, x, y));
      imRef(bg, x, y) = abs(50-imRef(I, x, y));
      if (imRef(fg, x, y ) > robustThreshold) 
	imRef(fg, x, y ) = robustThreshold;
      if (imRef(bg, x, y ) > robustThreshold) 
	imRef(bg, x, y ) = robustThreshold;
    }
  }
}


void loadModel(image<uchar> *I, int &K, 
	       matrix<double> * &horz, matrix<double> * &vert,
	       image<double> ** &data) {
  K = 4;
  horz = new matrix<double>(K,K);
  vert = new matrix<double>(K,K);
    /**/
  double lambda = 20;
  horz->init(lambda);
  vert->init(lambda);

  for (int i = 0; i < K; i++ ) {
    matRef(horz,i,i) = 0;
    matRef(vert,i,i) = 0;
  }
  
  int inMatHorz[4][4] = {	  {0,1000,5,100},
						  {1000,0,100,100},
						  {1000,100,0,5},
						  {5,100,1000,0}       };

  int inMatVert[4][4] = {	  {0,100,5,5},
						  {1000,0,1000,1000},
						  {1000,5,0,100},
						  {1000,5,100,0}       };

  for (int i_index = 0; i_index < K ; i_index++) {
	for (int j_index = 0; j_index < K ; j_index++) {
	  matRef(horz,i_index,j_index) = inMatHorz[i_index][j_index] * lambda;
	  matRef(vert,i_index,j_index) = inMatVert[i_index][j_index] * lambda;
	}
  }

  // data model
  image<double> *fg, *bg;
  loadData(I, fg, bg);
  data = new image<double> *[K];
  data[0] = bg;
  data[1] = bg;
  for (int i = 2; i < K; i++)
    data[i] = fg;
}

void image2outMat(int width, int height, vlib::image<double>* input, double *out) {
    for (int x=0; x < width; x++) {
        for (int y=0; y < height; y++) {
            out[y*width+x] = imRef(input, y, x);
        }
    }
}

void matrix2outMat(int width, int height, vlib::matrix<double>* input, double *out) {
    for (int x=0; x < width; x++) {
        for (int y=0; y < height; y++) {
            out[y*width+x] = imRef(input, y, x);
        }
    }
}

void mexFunction( int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) {
    // test args
    /* check for proper number of arguments */
    if(nrhs!=1) {
        mexErrMsgIdAndTxt("MyToolbox:arrayProduct:nrhs", "One inputs required.");
    }
    if(nlhs!=4) {
        mexErrMsgIdAndTxt("MyToolbox:arrayProduct:nlhs",  "Four outputs required.");
    }
    
    /* make sure the input argument is image */
    if( !mxIsDouble(prhs[0])) {
        mexErrMsgIdAndTxt("MyToolbox:arrayProduct:prhs", "Input must be uint8 matrix.");
    }
    
    // convert input to image
    int width = mxGetM (prhs[0]);
    int height = mxGetN (prhs[0]);
    vlib::image<vlib::uchar> *input = new vlib::image<vlib::uchar>(width, height);
    
    // insert data to image objec (equivalent to loadPGM()  )
    double *inMatrix;
    inMatrix = mxGetPr(prhs[0]);
    for (int x=0; x < width; x++) {
       for (int y=0; y < height; y++) {
           imRef(input, y, x) = inMatrix[y*width+x];
        }
    }
    
    // calling the logic
    int K = 4;
	 matrix<double> *horz, *vert;
	image<double> ** data;    
    loadModel(input, K, horz, vert, data);

    /* create the output matrix*/
    plhs[0] = mxCreateDoubleScalar(K);
    mwSize dataDims[1]; // creating dims array
    dataDims[0] = K; // we need only an array on K cells
    plhs[1] =mxCreateCellArray(1, dataDims);
    for (int i = 0; i<K; i++) {
        mxSetCell( plhs[1], i, mxCreateDoubleMatrix(height,width,mxREAL) );
    }
    plhs[2] = mxCreateDoubleMatrix(K,K,mxREAL);
    plhs[3] = mxCreateDoubleMatrix(K,K,mxREAL);
 
 
    /* get a pointer to the real data in the output matrix  */
   for (int i = 0; i<K; i++) {
        double *datatemp =  mxGetPr(mxGetCell( plhs[1], i));
        image2outMat(width, height,data[i],datatemp);
    } 
   
    double *out2 = mxGetPr(plhs[2]);
    matrix2outMat(K, K,horz,out2);

    double *out3 = mxGetPr(plhs[3]);
    matrix2outMat(K, K, vert, out3);
    
    // free memory
    delete data[0];
    delete data[2];
    delete data;
    delete horz;
    delete vert;
}